package study.dq.unit.test.until;


import org.dbunit.DatabaseUnitException;
import org.dbunit.database.DatabaseConnection;
import org.dbunit.database.IDatabaseConnection;
import org.dbunit.database.QueryDataSet;
import org.dbunit.dataset.DataSetException;
import org.dbunit.dataset.IDataSet;
import org.dbunit.dataset.xml.FlatXmlDataSet;
import org.dbunit.dataset.xml.FlatXmlProducer;
import org.dbunit.operation.DatabaseOperation;
import org.junit.AfterClass;
import org.junit.BeforeClass;
import org.xml.sax.InputSource;

import java.io.*;
import java.sql.SQLException;
import static org.junit.Assert.*;

public abstract class AbstractDbunitTestCase {

    public static IDatabaseConnection conn;
    private File tempFile;

    @BeforeClass
    public static void init() throws DatabaseUnitException {
        conn = new DatabaseConnection(DBUtil.getConnection());
    }

    protected IDataSet createDataSet(String tname) throws SQLException, DataSetException {
        InputStream is =  AbstractDbunitTestCase.class.getClassLoader().getResourceAsStream("dbunit_xml/"+tname+".xml");
        assertNotNull("dbunit的基本数据文件不存在",is);
        return new FlatXmlDataSet(new FlatXmlProducer(new InputSource(is)));
    }

    private void writeBackupFile(IDataSet dataSet) throws IOException, DataSetException {
        tempFile = File.createTempFile("backup","xml");
        System.out.println(tempFile.getAbsolutePath());
        FlatXmlDataSet.write(dataSet,new FileWriter(tempFile));
    }


    protected void backupAllTable() throws SQLException, IOException, DataSetException {
        IDataSet dataSet = conn.createDataSet();
        writeBackupFile(dataSet);
    }

    protected void resumeAllTable() throws DatabaseUnitException, SQLException, FileNotFoundException {
        IDataSet dataSet = new FlatXmlDataSet(new FlatXmlProducer(new InputSource(new FileInputStream(tempFile))));
        DatabaseOperation.CLEAN_INSERT.execute(conn, dataSet);
    }

    protected void backupCustomTable(String[] tname) throws DataSetException, IOException {
        QueryDataSet ds = new QueryDataSet(conn);
        for(String str:tname) {
            ds.addTable(str);
        }
        writeBackupFile(ds);
    }

    protected void bakcupOneTable(String tname) throws DataSetException, IOException {
        backupCustomTable(new String[]{tname});
    }

    @AfterClass
    public static void destory() throws SQLException {
        if (conn != null&& conn.getConnection()!=null) {
            try {
                DBUtil.closeConnection(conn.getConnection());
            } catch (SQLException e) {
                e.printStackTrace();
            }
        }
    }
}
